package com.example.automotive.util;

import android.graphics.Bitmap;
import android.graphics.Canvas;
import android.graphics.Color;
import android.graphics.Paint;
import android.util.Log;
import android.widget.ImageView;
import android.view.View;

import androidx.camera.core.ImageProxy;

import com.example.automotive.MainActivity;
import com.example.automotive.util.DataManager;

import ai.onnxruntime.*;

import java.io.File;
import java.nio.FloatBuffer;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Arrays;
import java.util.ArrayList;
import com.example.automotive.R;

public class ONNXManager {

    private static final String TAG = "ONNXManager";
    private final MainActivity mainActivity;
    private OrtEnvironment env;
    private OrtSession session;
    private DataManager dataManager;

    public ONNXManager(MainActivity mainActivity) {
        this.mainActivity = mainActivity;
        dataManager = new DataManager("UpdateModel", mainActivity);
        try {
            env = OrtEnvironment.getEnvironment();
            File modelFile = new File(mainActivity.getFilesDir(), Config.MODEL_NAME);
            Log.d(TAG, "Model file path: " + modelFile.getAbsolutePath());
            if (!modelFile.exists()) {
                Log.e(TAG, "Model file does not exist");
                return;
            }
            session = env.createSession(modelFile.getAbsolutePath());

            Map<String, NodeInfo> inputInfoMap = session.getInputInfo();
            for (Map.Entry<String, NodeInfo> entry : inputInfoMap.entrySet()) {
                Log.d(TAG, "Input Name: " + entry.getKey());
                Log.d(TAG, "Input Info: " + entry.getValue().getInfo().toString());
            }
        } catch (OrtException e) {
            Log.e(dataManager.TAG, "Failed to initialize ONNX Runtime or load model", e);
        }
    }

    public void runModel(Bitmap bitmap) {
        try {
            Map<String, NodeInfo> inputInfoMap = session.getInputInfo();
            NodeInfo inputInfo = inputInfoMap.get("images"); // 使用正确的输入节点名称
            if (inputInfo == null) {
                Log.e(TAG, "Input info is null");
                return;
            }
            TensorInfo tensorInfo = (TensorInfo) inputInfo.getInfo();
            if (tensorInfo.type != OnnxJavaType.FLOAT) {
                Log.e(TAG, "Model input is not of type FLOAT");
                return;
            }

            long[] inputShape = tensorInfo.getShape();
            if (inputShape == null || inputShape.length != 4) {
                Log.e(TAG, "Invalid input shape");
                return;
            }

            int inputHeight = (int) inputShape[2];
            int inputWidth = (int) inputShape[3];
            Bitmap resizedBitmap = Bitmap.createScaledBitmap(bitmap, inputWidth, inputHeight, false);

            float[] floatValues = convertBitmapToFloatArray(resizedBitmap);

            // 打印缓冲区信息
            Log.d(TAG, "Float array length: " + floatValues.length);
            Log.d(TAG, "Expected elements count: " + (inputShape[1] * inputShape[2] * inputShape[3]));

            // 验证浮点数组的长度是否匹配预期的元素数量
            int expectedSize = (int) (inputShape[1] * inputShape[2] * inputShape[3]);
            if (floatValues.length != expectedSize) {
                Log.e(TAG, "Float array does not contain the expected number of elements. Expected: " + expectedSize + ", Actual: " + floatValues.length);
                return;
            }

            try (OnnxTensor tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(floatValues), inputShape)) {
                Log.d(TAG, "OnnxTensor created successfully");
                OrtSession.Result result = session.run(Collections.singletonMap("images", tensor));

                // 处理三维输出
                Object outputObj = result.get(0).getValue();
                if (outputObj instanceof float[][][]) {
                    float[][][] output = (float[][][]) outputObj;
                    Log.d(TAG, "Output shape: [" + output.length + ", " + output[0].length + ", " + output[0][0].length + "]");
                    for (int i = 0; i < output.length; i++) {
                        for (int j = 0; j < output[i].length; j++) {
                            Log.d(TAG, "Detection " + j + ": " + Arrays.toString(output[i][j]));
                        }
                    }
                    Bitmap resultBitmap = drawDetectionsOnBitmap(bitmap, output);
                    displayBitmap(resultBitmap);
                } else {
                    Log.e(TAG, "Unexpected output type: " + outputObj.getClass().getSimpleName());
                }
            } catch (OrtException e) {
                Log.e(TAG, "Failed to create OnnxTensor or run model", e);
            }

        } catch (OrtException e) {
            Log.e(TAG, "Failed to run model", e);
        }
    }

    private float[] convertBitmapToFloatArray(Bitmap bitmap) {
        int width = bitmap.getWidth();
        int height = bitmap.getHeight();
        int pixelCount = width * height;

        float[] floatArray = new float[3 * pixelCount];
        int[] intValues = new int[pixelCount];
        bitmap.getPixels(intValues, 0, width, 0, 0, width, height);

        for (int i = 0; i < pixelCount; i++) {
            int pixelValue = intValues[i];
            floatArray[i * 3] = ((pixelValue >> 16) & 0xFF) / 255.0f;  // Red
            floatArray[i * 3 + 1] = ((pixelValue >> 8) & 0xFF) / 255.0f;   // Green
            floatArray[i * 3 + 2] = (pixelValue & 0xFF) / 255.0f;          // Blue
        }

        // 打印前10个像素的归一化值
        for (int i = 0; i < Math.min(10, pixelCount); i++) {
            Log.d(TAG, "Pixel " + i + ": R=" + floatArray[i * 3] + ", G=" + floatArray[i * 3 + 1] + ", B=" + floatArray[i * 3 + 2]);
        }

        return floatArray;
    }



    private Bitmap drawDetectionsOnBitmap(Bitmap bitmap, float[][][] detections) {
        Bitmap resultBitmap = bitmap.copy(Bitmap.Config.ARGB_8888, true);
        Canvas canvas = new Canvas(resultBitmap);
        Paint paint = new Paint();
        paint.setColor(Color.RED);
        paint.setStyle(Paint.Style.STROKE);
        paint.setStrokeWidth(5);

        int originalWidth = bitmap.getWidth();
        int originalHeight = bitmap.getHeight();

        for (float[][] detectionArray : detections) {
            for (float[] detection : detectionArray) {
                Log.d(TAG, "Detection: " + Arrays.toString(detection)); // 打印每个检测结果

                // 假设前四个值是检测框坐标（相对值），第五个值是置信度
                if (detection[4] > 0.5) { // 设置置信度阈值为0.5
                    float left = detection[0] * originalWidth;
                    float top = detection[1] * originalHeight;
                    float right = detection[2] * originalWidth;
                    float bottom = detection[3] * originalHeight;
                    Log.d(TAG, String.format("Converted Coordinates - left: %f, top: %f, right: %f, bottom: %f", left, top, right, bottom));
                    canvas.drawRect(left, top, right, bottom, paint);
                }
            }
        }
        return resultBitmap;
    }

    private void displayBitmap(Bitmap bitmap) {
        mainActivity.runOnUiThread(() -> {
            ImageView imageView = mainActivity.findViewById(R.id.resultImageView);
            imageView.setImageBitmap(bitmap);
            imageView.setVisibility(View.VISIBLE); // 设置 ImageView 可见
        });
    }
}
